{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7a96a3ab",
   "metadata": {},
   "source": [
    "# Direct Preference Optimization (DPO) Training with SageMaker\n",
    "\n",
    "This notebook demonstrates how to use the **DPOTrainer** to fine-tune large language models using Direct Preference Optimization (DPO). DPO is a technique that trains models to align with human preferences by learning from preference data without requiring a separate reward model.\n",
    "\n",
    "## What is DPO?\n",
    "\n",
    "Direct Preference Optimization (DPO) is a method for training language models to follow human preferences. Unlike traditional RLHF (Reinforcement Learning from Human Feedback), DPO directly optimizes the model using preference pairs without needing a reward model.\n",
    "\n",
    "**Key Benefits:**\n",
    "- Simpler than RLHF - no reward model required\n",
    "- More stable training process\n",
    "- Direct optimization on preference data\n",
    "- Works with LoRA for efficient fine-tuning\n",
    "\n",
    "## Workflow Overview\n",
    "\n",
    "1. **Prepare Preference Dataset**: Upload preference data in JSONL format\n",
    "2. **Register Dataset**: Create a SageMaker AI Registry dataset\n",
    "3. **Configure DPO Trainer**: Set up model, training parameters, and resources\n",
    "4. **Execute Training**: Run the DPO fine-tuning job\n",
    "5. **Track Results**: Monitor training with MLflow integration"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2446b6a5",
   "metadata": {},
   "source": [
    "## Step 1: Prepare and Register Preference Dataset\n",
    "\n",
    "DPO requires preference data in a specific format where each example contains:\n",
    "- **prompt**: The input text\n",
    "- **chosen**: The preferred response\n",
    "- **rejected**: The less preferred response\n",
    "\n",
    "The dataset should be in JSONL format with each line containing one preference example."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Configure AWS credentials and region\n",
    "#! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n",
    "#! aws configure set region us-west-2"
   ],
   "id": "3878997b198befc0"
  },
  {
   "cell_type": "code",
   "id": "ed5d2927f430664b",
   "metadata": {},
   "source": [
    "from sagemaker.ai_registry.dataset import DataSet\n",
    "from sagemaker.ai_registry.dataset_utils import CustomizationTechnique\n",
    "\n",
    "# Upload dataset to S3\n",
    "import boto3\n",
    "s3 = boto3.client('s3')\n",
    "s3.upload_file(\n",
    "    './dpo-preference_dataset_train_256.jsonl',\n",
    "    'nova-mlflow-us-west-2',\n",
    "    'dataset/preference_dataset_train_256.jsonl'\n",
    ")\n",
    "\n",
    "# Register dataset in SageMaker AI Registry\n",
    "# This creates a versioned dataset that can be referenced by ARN\n",
    "dataset = DataSet.create(\n",
    "    name=\"demo-6\",\n",
    "    data_location=\"s3://nova-mlflow-us-west-2/dataset/preference_dataset_train_256.jsonl\", \n",
    "    customization_technique=CustomizationTechnique.DPO, \n",
    "    wait=True\n",
    ")\n",
    "\n",
    "print(f\"Dataset ARN: {dataset.arn}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "71071d5c",
   "metadata": {},
   "source": [
    "## Step 2: Configure and Execute DPO Training\n",
    "\n",
    "The **DPOTrainer** provides a high-level interface for DPO fine-tuning with the following key features:\n",
    "\n",
    "### Key Parameters:\n",
    "- **model**: Base model to fine-tune (from SageMaker Hub)\n",
    "- **training_type**: Fine-tuning method (LoRA recommended for efficiency)\n",
    "- **training_dataset**: ARN of the registered preference dataset\n",
    "- **model_package_group_name**: Where to store the fine-tuned model\n",
    "- **mlflow_resource_arn**: MLflow tracking server for experiment logging\n",
    "\n",
    "### Training Features:\n",
    "- **Serverless Training**: Automatically managed compute resources\n",
    "- **LoRA Integration**: Parameter-efficient fine-tuning\n",
    "- **MLflow Tracking**: Automatic experiment and metrics logging\n",
    "- **Model Versioning**: Automatic model package creation"
   ]
  },
  {
   "cell_type": "code",
   "id": "e42719df1e792227",
   "metadata": {},
   "source": [
    "import random\n",
    "! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n",
    "! aws configure set region  us-west-2\n",
    "\n",
    "from sagemaker.train.dpo_trainer import DPOTrainer\n",
    "from sagemaker.train.common import TrainingType\n",
    "\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "0352bdaa-fa13-44c5-a70c-0d9bf7a10477",
   "metadata": {},
   "source": [
    "# Create DPOTrainer instance with comprehensive configuration\n",
    "trainer = DPOTrainer(\n",
    "    # Base model from SageMaker Hub\n",
    "    model=\"meta-textgeneration-llama-3-2-1b-instruct\",\n",
    "    \n",
    "    # Use LoRA for efficient fine-tuning\n",
    "    training_type=TrainingType.LORA,\n",
    "    \n",
    "    # Model versioning and storage\n",
    "    model_package_group_name=\"arn:aws:sagemaker:us-west-2:<>:model-package-group/test-finetuned-models-gamma\",\n",
    "    \n",
    "    # MLflow experiment tracking\n",
    "    mlflow_resource_arn=\"arn:aws:sagemaker:us-west-2:<>:mlflow-tracking-server/ashwpat-test\",\n",
    "    \n",
    "    # Training data (from Step 1)\n",
    "    training_dataset=\"arn:aws:sagemaker:us-west-2:<>:hub-content/F3LMYANDKWPZCROJVCKMJ7TOML6QMZBZRRQOVTUL45VUK7PJ4SXA/DataSet/demo-6/0.0.4\",\n",
    "    \n",
    "    # Output configuration\n",
    "    s3_output_path=\"s3://nova-mlflow-us-west-2/output\",\n",
    "    \n",
    "    # IAM role for training job\n",
    "    role=\"arn:aws:iam::<>:role/Admin\",\n",
    "    \n",
    "    # Unique job name\n",
    "    base_job_name=f\"dpo-llama-{random.randint(1, 1000)}\",\n",
    ")\n",
    "\n",
    "# Customize training hyperparameters\n",
    "# DPO-specific parameters are automatically loaded from the model's recipe\n",
    "trainer.hyperparameters.max_epochs = 1  # Quick training for demo\n",
    "\n",
    "print(\"Starting DPO training job...\")\n",
    "print(f\"Job name: {trainer.base_job_name}\")\n",
    "print(f\"Base model: {trainer._model_name}\")\n",
    "\n",
    "# Execute training with monitoring\n",
    "training_job = trainer.train(wait=True)\n",
    "\n",
    "print(f\"Training completed! Job ARN: {training_job.training_job_arn}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "22f6a210-0a0c-4b7a-af4d-2e08eae1c048",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "from sagemaker.core.utils.utils import Unassigned\n",
    "import json\n",
    "\n",
    "print(json.dumps({k: v for k, v in training_job.__dict__.items() if not isinstance(v, Unassigned) and \"Unassigned object\" not in str(v)}, indent=2, default=str))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "73d7545b",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "\n",
    "After training completes, you can:\n",
    "\n",
    "1. **Deploy the Model**: Use `ModelBuilder` to deploy the fine-tuned model\n",
    "2. **Evaluate Performance**: Compare responses from base vs fine-tuned model\n",
    "3. **Monitor Metrics**: Review training metrics in MLflow\n",
    "4. **Iterate**: Adjust hyperparameters and retrain if needed\n",
    "\n",
    "### Example Deployment:\n",
    "```python\n",
    "from sagemaker.serve import ModelBuilder\n",
    "\n",
    "# Deploy the fine-tuned model\n",
    "model_builder = ModelBuilder(model=training_job)\n",
    "model_builder.build(role_arn=\"arn:aws:iam::account:role/SageMakerRole\")\n",
    "endpoint = model_builder.deploy(endpoint_name=\"dpo-finetuned-llama\")\n",
    "```\n",
    "\n",
    "The fine-tuned model will now generate responses that better align with the preferences in your training data."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
